package grpc_proxy_middleware

import (
	"encoding/json"
	"errors"
	"fmt"
	"gitee.com/zhoulvvv/my_go_gateway/dao"
	"gitee.com/zhoulvvv/my_go_gateway/public"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/peer"
	"log"
	"strings"
)

// GRPC 租户限流
func GRPCJwtFlowLimitModeMiddleware(service *dao.ServiceDetail) func(srv interface{}, ss grpc.ServerStream,
	info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		// 获取租户信息
		md, ok := metadata.FromIncomingContext(ss.Context())
		if !ok {
			log.Printf("【GRPC 租户限流】 failed with error metadata.FromIncomingContext(ss.Context())")
		}
		// 获取租户信息
		apps := md.Get("app")
		if len(apps) == 0 {
			if err := handler(srv, ss); err != nil {
				log.Printf("【GRPC】RPC failed with error %v\n", err)
			}
			return nil
		}
		app := &dao.App{}
		if err := json.Unmarshal([]byte(apps[0]), app); err != nil {
			log.Printf("【GRPC 租户限流】 json.Unmarshal([]byte(apps[0]),app) %v\n", err)
		}
		appQps := float64(app.Qps)
		// 如果租户限流QPS是0就表示无限制
		if appQps > 0 {
			peerCtx, ok := peer.FromContext(ss.Context())
			if !ok {
				log.Printf("【GRPC IP 白名单中间件中获取客户端IP失败】peer.FromContext(ss.Context()) failed with error")
			}
			splits := strings.Split(peerCtx.Addr.String(), ":")
			// log.Printf("peerCtx.Addr.String():%s\n", peerCtx.Addr.String())
			clientIP := ""
			if len(splits) == 2 {
				clientIP = splits[0]
			}
			// 租户限流，根据租户下不同的访问IP进行限流
			serverLimiter, err := public.FlowLimiterHandler.GetLimiter(
				public.FlowAppPrefix+app.AppId+"_"+clientIP, appQps)
			if err != nil {
				return errors.New("【GRPC 租户限流】获取租户限流器失败：" + err.Error())
			}
			if !serverLimiter.Allow() {
				return errors.New(fmt.Sprintf("【GRPC 租户限流】租户限流中，每秒最大QPS：%v，限流IP：%s", appQps, clientIP))
			}
		}
		if err := handler(srv, ss); err != nil {
			log.Printf("【GRPC】RPC failed with error %v\n", err)
		}
		return nil
	}
}
